# 用法python panorama.py datasets/1
import cv2
import numpy as np
import os
import time
import argparse
from typing import List, Tuple

def register_images(images: List[np.ndarray]) -> np.ndarray:
    """图像配准与拼接"""
    if len(images) < 2:
        return images[0] if images else None
    
    start_time = time.time()
    
    # 使用SIFT特征检测器（比ORB更稳定）
    sift = cv2.SIFT_create(nfeatures=5000, contrastThreshold=0.04, edgeThreshold=10)
    
    # 预处理所有图像并提取特征
    # print("提取特征点...")
    gray_images = []
    keypoints = []
    descriptors = []
    
    for i, img in enumerate(images):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        gray = cv2.medianBlur(gray, 3)
        gray_images.append(gray)
        
        kp, des = sift.detectAndCompute(gray, None)
        keypoints.append(kp)
        descriptors.append(des)
        
        progress = (i + 1) / len(images) * 100
        print(f"\r特征提取进度: [{get_progress_bar(progress)}] {progress:.1f}%", end="", flush=True)
    print()
    
    # 构建图像连接图
    # print("构建图像连接...")
    connections = build_image_connections(images, keypoints, descriptors)
    
    if not connections:
        print("错误: 无法找到足够的图像连接")
        return simple_stitch(images)
    
    # 找到中心图像
    center_idx = find_center_image(connections, len(images))
    #print(f"使用图像 {center_idx + 1} 作为中心参考")
    
    # 计算相对于中心图像的变换矩阵
    #print("计算变换矩阵...")
    transforms = compute_transforms(center_idx, images, keypoints, descriptors, connections)
    
    if len(transforms) < 2:
        print("错误: 无法计算足够的变换矩阵")
        return simple_stitch(images)
    
    # 优化变换矩阵
    transforms = optimize_transforms(transforms, connections, center_idx)
    
    # 拼接所有图像
    # print("开始图像拼接...")
    result = stitch_all_images(images, transforms, center_idx)
    
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"拼接完成，总耗时: {elapsed_time:.2f}秒")
    
    return result

def simple_stitch(images):
    """简单的拼接方法，用于备选"""
    #print("使用简单拼接方法...")
    stitcher = cv2.Stitcher.create(cv2.Stitcher_PANORAMA)
    status, result = stitcher.stitch(images)
    
    if status == cv2.Stitcher_OK:
        return result
    else:
        # 如果OpenCV拼接失败，返回第一张图像
        return images[0]

def build_image_connections(images, keypoints, descriptors, min_matches=20):
    """构建图像之间的连接关系"""
    connections = []
    # 使用FLANN匹配器
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    for i in range(len(images)):
        for j in range(i + 1, len(images)):
            if (descriptors[i] is None or descriptors[j] is None or 
                len(descriptors[i]) < min_matches or len(descriptors[j]) < min_matches):
                continue
            
            try:
                # 特征匹配
                matches = flann.knnMatch(descriptors[i], descriptors[j], k=2)
                good_matches = []
                for m, n in matches:
                    if m.distance < 0.7 * n.distance:
                        good_matches.append(m)
                
                if len(good_matches) >= min_matches:
                    connections.append((i, j, len(good_matches)))
                    #print(f"图像 {i+1} 和图像 {j+1} 找到 {len(good_matches)} 个匹配点")
            except Exception as e:
                print(f"匹配图像 {i+1} 和 {j+1} 时出错: {e}")
                continue
    
    return connections

def find_center_image(connections, num_images):
    """找到连接最多的图像作为中心"""
    if not connections:
        return num_images // 2
    
    connection_count = [0] * num_images
    for i, j, count in connections:
        connection_count[i] += count
        connection_count[j] += count
    
    return np.argmax(connection_count)

def compute_transforms(center_idx, images, keypoints, descriptors, connections):
    """计算所有图像相对于中心图像的变换矩阵"""
    transforms = {center_idx: np.eye(3, dtype=np.float32)}
    queue = [center_idx]
    visited = set([center_idx])
    
    # 使用FLANN匹配器
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    while queue:
        current = queue.pop(0)
        
        # 找到与当前图像连接的其他图像
        for i, j, match_count in connections:
            neighbor = None
            if i == current and j not in visited:
                neighbor = j
                src_idx, dst_idx = j, current
            elif j == current and i not in visited:
                neighbor = i
                src_idx, dst_idx = i, current
            
            if neighbor is not None and descriptors[src_idx] is not None and descriptors[dst_idx] is not None:
                try:
                    # 特征匹配
                    matches = flann.knnMatch(descriptors[src_idx], descriptors[dst_idx], k=2)
                    good_matches = []
                    for m, n in matches:
                        if m.distance < 0.7 * n.distance:
                            good_matches.append(m)
                    
                    if len(good_matches) >= 10:
                        # 计算单应矩阵
                        src_pts = np.float32([keypoints[src_idx][m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
                        dst_pts = np.float32([keypoints[dst_idx][m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
                        
                        # 使用更严格的RANSAC参数
                        H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 
                                                    ransacReprojThreshold=2.0, 
                                                    maxIters=3000, 
                                                    confidence=0.999)
                        
                        if H is not None and is_homography_valid(H):
                            # 计算相对于中心图像的变换
                            transforms[src_idx] = transforms[dst_idx] @ H
                            queue.append(src_idx)
                            visited.add(src_idx)
                            # print(f"图像 {src_idx + 1} -> 图像 {dst_idx + 1} 配准成功 (内点: {np.sum(mask)}/{len(good_matches)})")
                except Exception as e:
                    print(f"计算变换时出错: {e}")
                    continue
    
    return transforms

def optimize_transforms(transforms, connections, center_idx):
    """优化变换矩阵以减少累积误差"""
    #print("优化变换矩阵...")
    
    # 这里可以添加束调整等优化算法
    # 目前简单地对变换矩阵进行规范化
    for idx, H in transforms.items():
        if idx != center_idx:
            # 确保变换矩阵的尺度合理
            det = np.linalg.det(H)
            if abs(det) > 1e-6:
                # 轻微调整以避免极端变换
                H_normalized = H / np.cbrt(det)  # 使用立方根保持方向
                transforms[idx] = H_normalized
    
    return transforms

def stitch_all_images(images, transforms, center_idx):
    """使用变换矩阵拼接所有图像"""
    # 计算画布尺寸
    all_corners = []
    valid_transforms = []
    
    for idx, H in transforms.items():
        h, w = images[idx].shape[:2]
        corners = cv2.perspectiveTransform(
            np.float32([[0,0], [0,h], [w,h], [w,0]]).reshape(-1,1,2),
            H
        )
        all_corners.append(corners)
        valid_transforms.append((idx, H))
    
    all_corners = np.concatenate(all_corners)
    x_min = int(np.floor(np.min(all_corners[..., 0])))
    y_min = int(np.floor(np.min(all_corners[..., 1])))
    x_max = int(np.ceil(np.max(all_corners[..., 0])))
    y_max = int(np.ceil(np.max(all_corners[..., 1])))
    
    # 限制最大画布尺寸，避免内存溢出
    max_canvas_size = 10000
    if (x_max - x_min) > max_canvas_size or (y_max - y_min) > max_canvas_size:
        scale = max_canvas_size / max(x_max - x_min, y_max - y_min)
        x_min = int(x_min * scale)
        y_min = int(y_min * scale)
        x_max = int(x_max * scale)
        y_max = int(y_max * scale)
        print(f"画布尺寸过大，缩放系数: {scale:.2f}")
    
    H_translation = np.array([
        [1, 0, -x_min], 
        [0, 1, -y_min], 
        [0, 0, 1]
    ], dtype=np.float32)
    
    canvas_width = x_max - x_min
    canvas_height = y_max - y_min
    
    #print(f"画布尺寸: {canvas_width} x {canvas_height}")
    
    result = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8)
    result_mask = np.zeros((canvas_height, canvas_width), dtype=np.uint8)
    
    # 按连接顺序融合图像（中心图像优先）
    #print("图像融合...")
    sorted_indices = sorted(transforms.keys(), key=lambda x: (abs(x - center_idx), x))
    
    for i, idx in enumerate(sorted_indices):
        H_total = H_translation @ transforms[idx]
        
        # 使用双线性插值进行变换，避免放射状模糊
        img_warped = cv2.warpPerspective(
            images[idx],
            H_total,
            (canvas_width, canvas_height),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(0,0,0)
        )
        
        # 创建精确的掩码
        mask_warped = cv2.warpPerspective(
            np.ones((images[idx].shape[0], images[idx].shape[1]), dtype=np.uint8) * 255,
            H_total,
            (canvas_width, canvas_height),
            flags=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=0
        )
        
        # 使用改进的融合方法
        if i == 0:
            result = img_warped.copy()
            result_mask = mask_warped
        else:
            result = improved_blending(result, img_warped, result_mask, mask_warped)
            result_mask = cv2.bitwise_or(result_mask, mask_warped)
        
        progress = (i + 1) / len(sorted_indices) * 100
        print(f"\r融合进度: [{get_progress_bar(progress)}] {progress:.1f}%", end="", flush=True)
    
    print()
    return result

# def improved_blending(img1, img2, mask1, mask2):
#     """改进的图像融合方法，避免放射状模糊"""
#     # 创建结果图像
#     result = img1.copy()
    
#     # 找到重叠区域
#     overlap = cv2.bitwise_and(mask1, mask2)
    
#     if not np.any(overlap):
#         # 如果没有重叠，直接复制第二张图像
#         result[mask2 > 0] = img2[mask2 > 0]
#         return result
    
#     # 使用更精确的掩码处理
#     mask1_only = cv2.bitwise_and(mask1, cv2.bitwise_not(overlap))
#     mask2_only = cv2.bitwise_and(mask2, cv2.bitwise_not(overlap))
    
#     # 非重叠区域直接复制
#     result[mask2_only > 0] = img2[mask2_only > 0]
    
#     # 重叠区域使用线性混合
#     if np.any(overlap):
#         # 创建距离变换权重图
#         dist1 = cv2.distanceTransform(mask1, cv2.DIST_L2, 5)
#         dist2 = cv2.distanceTransform(mask2, cv2.DIST_L2, 5)
        
#         # 对重叠区域进行高斯模糊
#         overlap_float = overlap.astype(np.float32) / 255.0
#         overlap_blur = cv2.GaussianBlur(overlap_float, (51, 51), 20)
        
#         # 计算权重
#         weight1 = overlap_blur
#         weight2 = 1.0 - overlap_blur
        
#         # 扩展权重到3通道
#         weight1_3ch = np.repeat(weight1[:, :, np.newaxis], 3, axis=2)
#         weight2_3ch = np.repeat(weight2[:, :, np.newaxis], 3, axis=2)
        
#         # 转换为浮点数进行混合
#         img1_float = img1.astype(np.float32)
#         img2_float = img2.astype(np.float32)
        
#         # 在重叠区域进行混合
#         blended = img1_float * weight1_3ch + img2_float * weight2_3ch
        
#         # 只更新重叠区域
#         overlap_3ch = np.repeat(overlap[:, :, np.newaxis], 3, axis=2)
#         result = result.astype(np.float32)
#         result[overlap_3ch > 0] = blended[overlap_3ch > 0]
#         result = np.clip(result, 0, 255).astype(np.uint8)
    
#     return result
def improved_blending(img1, img2, mask1, mask2):
    """改进的图像融合方法，避免黑色缝纫线"""
    # 平滑掩码边界，避免硬线
    kernel = np.ones((5,5), np.uint8)
    mask1_smooth = cv2.erode(mask1, kernel, iterations=1)
    mask2_smooth = cv2.erode(mask2, kernel, iterations=1)

    result = img1.copy()

    # 重叠区域
    overlap = cv2.bitwise_and(mask1_smooth, mask2_smooth)

    # 非重叠区域
    result[mask2_smooth > 0] = img2[mask2_smooth > 0]

    if np.any(overlap):
        dist1 = cv2.distanceTransform(mask1_smooth, cv2.DIST_L2, 5)
        dist2 = cv2.distanceTransform(mask2_smooth, cv2.DIST_L2, 5)

        weight1 = dist1 / (dist1 + dist2 + 1e-6)
        weight2 = 1.0 - weight1

        weight1_3 = np.repeat(weight1[:, :, None], 3, axis=2)
        weight2_3 = np.repeat(weight2[:, :, None], 3, axis=2)

        blended = img1.astype(np.float32) * weight1_3 + img2.astype(np.float32) * weight2_3
        mask_overlap_3 = np.repeat(overlap[:, :, None], 3, axis=2)

        result = result.astype(np.float32)
        result[mask_overlap_3 > 0] = blended[mask_overlap_3 > 0]
        result = result.astype(np.uint8)

    return result

def get_progress_bar(progress, length=20):
    """生成进度条"""
    filled = int(length * progress / 100)
    bar = '█' * filled + '░' * (length - filled)
    return bar

def is_homography_valid(H):
    """更严格的单应性矩阵验证"""
    if H is None:
        return False
    
    # 检查行列式，避免过度缩放
    det = np.linalg.det(H)
    if abs(det) < 1e-8 or abs(det) > 100:
        # print(f"无效行列式: {det}")
        return False
    
    # 检查变换后的角点是否合理
    test_points = np.float32([[0,0], [100,0], [100,100], [0,100]]).reshape(-1,1,2)
    transformed = cv2.perspectiveTransform(test_points, H)
    
    # 检查变换后是否出现严重扭曲
    area_original = 100 * 100
    area_transformed = cv2.contourArea(transformed)
    area_ratio = area_transformed / area_original
    
    if area_ratio < 0.01 or area_ratio > 100:
        # print(f"无效面积比: {area_ratio}")
        return False
    
    # 检查角点顺序是否保持（避免镜像变换）
    if not is_convex_quadrilateral(transformed):
        print("角点顺序异常")
        return False
    
    return True

def is_convex_quadrilateral(points):
    """检查四点是否构成凸四边形"""
    points = points.reshape(4, 2)
    
    # 计算叉积检查凸性
    def cross(o, a, b):
        return (a[0]-o[0])*(b[1]-o[1]) - (a[1]-o[1])*(b[0]-o[0])
    
    # 检查所有顶点是否按相同方向排列
    signs = []
    for i in range(4):
        o = points[i]
        a = points[(i+1)%4]
        b = points[(i+2)%4]
        signs.append(cross(o, a, b) > 0)
    
    return all(signs) or not any(signs)

def remove_black_borders(image: np.ndarray, threshold: int = 10) -> np.ndarray:
    """移除黑边"""
    if image is None or image.size == 0:
        return image
        
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)
    
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return image
    
    largest_contour = max(contours, key=lambda c: cv2.contourArea(c))
    x, y, w, h = cv2.boundingRect(largest_contour)
    
    # 确保裁剪区域有效
    if w > 50 and h > 50:  # 最小尺寸限制
        cropped = image[y:y+h, x:x+w]
        return cropped
    else:
        return image

# def load_images_from_folder(folder: str) -> List[np.ndarray]:
#     """加载彩色图片"""
#     images = []
#     valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif')
#
#     for filename in sorted(os.listdir(folder)):
#         if filename.lower().endswith(valid_extensions):
#             img_path = os.path.join(folder, filename)
#             img = cv2.imread(img_path, cv2.IMREAD_COLOR)
#             if img is not None:
#                 # 调整图像尺寸，避免过大
#                 h, w = img.shape[:2]
#                 max_dimension = 1200  # 降低分辨率以减少模糊
#                 if max(h, w) > max_dimension:
#                     scale = max_dimension / max(h, w)
#                     new_w = int(w * scale)
#                     new_h = int(h * scale)
#                     img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
#                 images.append(img)
#                 # print(f"已加载: {filename} ({img.shape[1]}x{img.shape[0]})")
#     return images

def load_images_from_paths(image_paths):
    """从图像路径列表加载图像"""
    images = []
    valid_paths = []

    for path in image_paths:
        # 检查文件是否存在
        if not os.path.exists(path):
            print(f"警告: 文件 '{path}' 不存在，已跳过")
            continue

        # 检查是否为图像文件
        if not path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif')):
            print(f"警告: 文件 '{path}' 不是支持的图像格式，已跳过")
            continue

        # 读取图像
        img = cv2.imread(path, cv2.IMREAD_COLOR)
        if img is not None:
            # 调整图像尺寸，避免过大
            h, w = img.shape[:2]
            max_dimension = 1200  # 降低分辨率以减少模糊
            if max(h, w) > max_dimension:
                scale = max_dimension / max(h, w)
                new_w = int(w * scale)
                new_h = int(h * scale)
                img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
            images.append(img)
        if img is None:
            print(f"警告: 无法读取图像 '{path}'，已跳过")
            continue

        # images.append(img)
        valid_paths.append(path)

    if len(valid_paths) > 0:
        # print(f"成功加载 {len(valid_paths)} 张图片:")
        for path in valid_paths:
            print(f"  - {path}")
    else:
        print("未成功加载任何图片")

    return images

def stitch(paths):
    """主函数，处理命令行参数"""
    # parser = argparse.ArgumentParser(description='图像拼接工具')
    # parser.add_argument('C:/Users/dxj/Desktop/CV/datasets/航拍_平移/1', help='包含待拼接图片的文件夹路径')
    #
    # args = parser.parse_args()
    #
    # # 检查文件夹是否存在
    # if not os.path.exists(args.image_folder):
    #     print(f"错误: 文件夹 '{args.image_folder}' 不存在")
    #     return
    #
    # if not os.path.isdir(args.image_folder):
    #     print(f"错误: '{args.image_folder}' 不是一个文件夹")
    #     return
    
    images = load_images_from_paths(paths)
    
    if len(images) < 2:
        print("需要至少2张图片进行拼接")
        if len(images) == 1:
            cv2.imwrite("result.jpg", images[0])
            print("已保存单张图片为 result.jpg")
        return False, None
    
    # print(f"已加载 {len(images)} 张图片")
    
    try:
        stitched = register_images(images)
        if stitched is not None:
            final_result = remove_black_borders(stitched, threshold=15)
            return True, final_result
        else:
            print("拼接失败")
            return False, None

    except Exception as e:
        print(f"处理过程出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return False, None

if __name__ == "__main__":
    paths = ['./datasets/hangpai/1/DSC00378.JPG', './datasets/hangpai/1/DSC00379.JPG', './datasets/hangpai/1/DSC00380.JPG', './datasets/hangpai/1/DSC00381.JPG', './datasets/hangpai/1/DSC00382.JPG', './datasets/hangpai/1/DSC00383.JPG', './datasets/hangpai/1/DSC00384.JPG']
    stitch(paths)